package com.wfh.aihelper.ai.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * @Title: RagConfig
 * @Author wangfenghuan
 * @Package com.wfh.aihelper.ai.rag
 * @Date 2025/7/10 18:34
 * @description: RAG配置
 */
@Configuration
public class RagConfig {

    @Resource
    private EmbeddingModel qwenEmbeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Bean
    public ContentRetriever contentRetriever() {
        // 1、加载文档
        List<Document> documentList = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        // 2、文档切割, 每一个文档按照段落进行分割，最大是1000个自负，每一次最多的是200个字符
        DocumentByParagraphSplitter documentByParagraphSplitter =
                new DocumentByParagraphSplitter(1000, 200);
        // 3、自定义文档加载起，把文档转换成向量并保存在内存中存储
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                // 文档切割器
                .documentSplitter(documentByParagraphSplitter)
                // 为了提高文档的质量，为每一个切割后的文档碎片添加文档名称作为元信息
                .textSegmentTransformer(textSegment -> TextSegment.from(textSegment.metadata().getString("file_name")
                        + "\n" + textSegment.text(), textSegment.metadata()))
                // 使用的向量模型
                .embeddingModel(qwenEmbeddingModel)
                // 向量存储方式（基于内存）
                .embeddingStore(embeddingStore)
                .build();
        // 4、加载文档
        ingestor.ingest(documentList);
        // 自定义内容加载起
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                // 最多五条结果
                .maxResults(5)
                // 过滤掉结果小雨0.75的结果
                .minScore(0.75)
                .build();

    }
}
